// Package storage provides compression utilities for storage backends
package storage

import (
	"fmt"

	"github.com/klauspost/compress/zstd"
	"github.com/pierrec/lz4/v4"
)

// CompressionType represents the type of compression to use
type CompressionType int

const (
	NoCompression CompressionType = iota
	LZ4Compression
	ZstdCompression
)

// CompressionConfig represents compression configuration
type CompressionConfig struct {
	Type            CompressionType
	Level           int
	EnableThreshold int // Minimum size to enable compression (bytes)
}

// DefaultCompressionConfig returns default compression configuration
func DefaultCompressionConfig() CompressionConfig {
	return CompressionConfig{
		Type:            LZ4Compression,
		Level:           0, // Default level
		EnableThreshold: 1024, // 1KB
	}
}

// Compressor handles compression and decompression
type Compressor struct {
	config CompressionConfig
	zstdEncoder *zstd.Encoder
	zstdDecoder *zstd.Decoder
}

// NewCompressor creates a new compressor
func NewCompressor(config CompressionConfig) (*Compressor, error) {
	c := &Compressor{
		config: config,
	}

	// Initialize zstd encoder/decoder if needed
	if config.Type == ZstdCompression {
		var err error
		c.zstdEncoder, err = zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(config.Level)))
		if err != nil {
			return nil, fmt.Errorf("failed to create zstd encoder: %w", err)
		}

		c.zstdDecoder, err = zstd.NewReader(nil)
		if err != nil {
			return nil, fmt.Errorf("failed to create zstd decoder: %w", err)
		}
	}

	return c, nil
}

// Compress compresses data using the configured compression type
func (c *Compressor) Compress(data []byte) ([]byte, error) {
	if len(data) < c.config.EnableThreshold {
		// Data too small, don't compress
		result := make([]byte, len(data)+1)
		result[0] = byte(NoCompression)
		copy(result[1:], data)
		return result, nil
	}

	switch c.config.Type {
	case LZ4Compression:
		return c.compressLZ4(data)
	case ZstdCompression:
		return c.compressZstd(data)
	default:
		// No compression
		result := make([]byte, len(data)+1)
		result[0] = byte(NoCompression)
		copy(result[1:], data)
		return result, nil
	}
}

// Decompress decompresses data
func (c *Compressor) Decompress(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("empty data to decompress")
	}

	compressionType := CompressionType(data[0])
	dataToDecompress := data[1:]

	switch compressionType {
	case NoCompression:
		return dataToDecompress, nil
	case LZ4Compression:
		return c.decompressLZ4(dataToDecompress)
	case ZstdCompression:
		return c.decompressZstd(dataToDecompress)
	default:
		return nil, fmt.Errorf("unknown compression type: %d", compressionType)
	}
}

func (c *Compressor) compressLZ4(data []byte) ([]byte, error) {
	compressed := make([]byte, lz4.CompressBlockBound(len(data)))
	n, err := lz4.CompressBlock(data, compressed, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to compress data with LZ4: %w", err)
	}

	// If compression didn't reduce size, store uncompressed
	if n >= len(data) {
		result := make([]byte, len(data)+1)
		result[0] = byte(NoCompression)
		copy(result[1:], data)
		return result, nil
	}

	// Store compressed data
	result := make([]byte, n+1)
	result[0] = byte(LZ4Compression)
	copy(result[1:], compressed[:n])
	return result, nil
}

func (c *Compressor) decompressLZ4(data []byte) ([]byte, error) {
	// Estimate decompressed size (usually 10x is safe for LZ4)
	decompressed := make([]byte, len(data)*10)
	n, err := lz4.UncompressBlock(data, decompressed)
	if err != nil {
		return nil, fmt.Errorf("failed to decompress data with LZ4: %w", err)
	}
	return decompressed[:n], nil
}

func (c *Compressor) compressZstd(data []byte) ([]byte, error) {
	compressed := c.zstdEncoder.EncodeAll(data, nil)

	// If compression didn't reduce size, store uncompressed
	if len(compressed) >= len(data) {
		result := make([]byte, len(data)+1)
		result[0] = byte(NoCompression)
		copy(result[1:], data)
		return result, nil
	}

	// Store compressed data
	result := make([]byte, len(compressed)+1)
	result[0] = byte(ZstdCompression)
	copy(result[1:], compressed)
	return result, nil
}

func (c *Compressor) decompressZstd(data []byte) ([]byte, error) {
	decompressed, err := c.zstdDecoder.DecodeAll(data, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to decompress data with zstd: %w", err)
	}
	return decompressed, nil
}

// Close closes the compressor and releases resources
func (c *Compressor) Close() error {
	if c.zstdEncoder != nil {
		c.zstdEncoder.Close()
	}
	if c.zstdDecoder != nil {
		c.zstdDecoder.Close()
	}
	return nil
}

// GetCompressionRatio calculates the compression ratio
func GetCompressionRatio(original, compressed int) float64 {
	if original == 0 {
		return 0
	}
	return float64(compressed) / float64(original)
}

// EstimateCompressedSize estimates the compressed size for given data
func EstimateCompressedSize(dataSize int, compressionType CompressionType) int {
	switch compressionType {
	case LZ4Compression:
		// LZ4 typically achieves 2-3x compression for text/data
		return dataSize / 2
	case ZstdCompression:
		// Zstd typically achieves 3-4x compression
		return dataSize / 3
	default:
		return dataSize
	}
}